#!/usr/bin/env python
# -*- coding: utf8 -*-
import socket, re
from prometheus_client import generate_latest, Gauge, Info
from prometheus_client.core import CollectorRegistry
from psutil import virtual_memory
from psutil import cpu_times


def check_port(ip, port):
    '''socket检测端口连通性'''
    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    s.settimeout(2)
    try:
        s.connect((ip, int(port)))
        s.shutdown(2)
        return True
    except:
        return False


REGISTRY = CollectorRegistry(auto_describe=False)
allportdata = [{'ip': 'xxx.xxx.xxx.xxx', 'port': 22}, {'ip': 'yyy.yyy.yyy.yyy', 'port': 22},
               {'ip': 'zzz.zzz.zzz.zzz', 'port': 22}]
mem_percent = Gauge(
    "system_memory_percent",
    "Total system memory percent.",
    registry=REGISTRY)
cpu_per = Gauge(
    "system_cpu_percent",
    "Total system cpu percent.",
    registry=REGISTRY)

mem_percent.set(virtual_memory().percent)
cpu_per.set(cpu_times().system)

METRIC_NAME_RE = re.compile(r'^[a-zA-Z_:][a-zA-Z0-9_:]*$')

name = 'ip_port_conn_status'
documentation = "IP PORT SOCKET CONNECT STATUS."
portcheck = Gauge(name=name, documentation=documentation, labelnames=('ip', 'port'), registry=REGISTRY)

for i in allportdata:
    status = check_port(i['ip'], i['port'])
    portcheck.labels(i['ip'], i['port']).inc(status)

result = generate_latest(REGISTRY).decode()
print(result)

